#!/usr/bin/env python3
"""
简化的ChromaDB测试
"""

import os
import sys
import asyncio

# 添加src目录到Python路径
src_path = os.path.join(os.path.dirname(__file__), 'src')
sys.path.insert(0, src_path)
os.environ['PYTHONPATH'] = src_path

try:
    import chromadb
    from chromadb.config import Settings
    print("ChromaDB导入成功")
except ImportError as e:
    print(f"ChromaDB导入失败: {e}")
    sys.exit(1)

def test_chromadb_direct():
    """直接测试ChromaDB"""
    print("=== 直接测试ChromaDB ===")
    
    try:
        # 使用配置文件中的路径
        persist_directory = "./data/chroma_db"
        abs_path = os.path.abspath(persist_directory)
        print(f"使用路径: {abs_path}")
        print(f"路径存在: {os.path.exists(abs_path)}")
        
        # 创建ChromaDB客户端
        client = chromadb.PersistentClient(
            path=persist_directory,
            settings=Settings(anonymized_telemetry=False)
        )
        print("ChromaDB客户端创建成功")
        
        # 获取conversations集合
        try:
            conversations_collection = client.get_collection('conversations')
            count = conversations_collection.count()
            print(f"conversations集合文档数量: {count}")
            
            if count > 0:
                # 获取前几个文档
                results = conversations_collection.get(limit=3)
                print(f"前3个文档ID: {results.get('ids', [])}")
                
                # 尝试搜索
                search_results = conversations_collection.query(
                    query_texts=["test"],
                    n_results=3
                )
                print(f"搜索'test'的结果数量: {len(search_results.get('ids', [[]])[0])}")
                
                # 尝试搜索中文
                search_results_cn = conversations_collection.query(
                    query_texts=["测试"],
                    n_results=3
                )
                print(f"搜索'测试'的结果数量: {len(search_results_cn.get('ids', [[]])[0])}")
                
        except Exception as e:
            print(f"获取conversations集合失败: {e}")
            # 列出所有集合
            collections = client.list_collections()
            print(f"可用集合: {[c.name for c in collections]}")
        
    except Exception as e:
        print(f"ChromaDB测试失败: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    test_chromadb_direct()